#!/usr/bin/env python
"""

    merge_insert_sizes.py
    [--log_file PATH]
    [--verbose]

"""

################################################################################
#
#   merge_insert_sizes
#
#
#   Copyright (c) 8/4/2011 Leo Goodstadt
#
#   Permission is hereby granted, free of charge, to any person obtaining a copy
#   of this software and associated documentation files (the "Software"), to deal
#   in the Software without restriction, including without limitation the rights
#   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#   copies of the Software, and to permit persons to whom the Software is
#   furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
#   THE SOFTWARE.
#################################################################################

import sys, os

# add self to search path for testing
if __name__ == '__main__':
    exe_path = os.path.split(os.path.abspath(sys.argv[0]))[0]
    module_name = os.path.split(sys.argv[0])[1]
    module_name = os.path.splitext(module_name)[0];
else:
    module_name = __name__

# Use import path from <<../python_modules>>
if __name__ == '__main__':
    sys.path.append("/home/lg/python_modules")


default_working_dir = "/data/mus/lg/projects/de_novo_assembly/working"



#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   options


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888


if __name__ == '__main__':
    import argparse
    from common_parser import get_common_parser

    parser = argparse.ArgumentParser(description='WHAT DOES THIS DO?',
                                    #formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                    parents=[get_common_parser()],
                                     )
    required_options = argparse._ArgumentGroup(parser, 'mandatory arguments')

    #
    #   Multiple input files
    #
    parser.add_argument("-i", "--input_file", dest="input_files",
                                    nargs='+',
                                    metavar="FILE",
                                    type= str,
                                    action="append",
                                  help="Name and path of input file.")
    #parser.add_argument("-r", "--output_r_file", dest="output_r_file",
    #                                metavar="FILE",
    #                                type= str,
    #                              help="R commands.")
    parser.add_argument("-p", "--output_pdf_file", dest="output_pdf_file",
                                    metavar="FILE",
                                    type= str,
                                  help="pdf file.")
    parser.set_defaults(
                        output_pdf_file   = os.path.join(default_working_dir, "insert_sizes.pdf"),                        )



    options = parser.parse_args()

    if not options.input_files:
        options.input_files =  [[os.path.join(default_working_dir, "*/*.insert_sizes_txt")
                                            ]]


    if not options.log_file:
        options.log_file            = os.path.join(default_working_dir, "merge_insert_sizes.log")


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   imports


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#from json import dumps
from collections import defaultdict
from subprocess import Popen, PIPE



#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

def array_in_r_text_format (name, data):
    """
    Returns list as a string for specifying vectors in R
    """
    return name + "=c(" + ",".join(map(str, data)) + ")"


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Logger


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

from lg_program_logging import MESSAGE, getLogger

if __name__ == '__main__':

    #
    #   set up log
    #
    logger = getLogger(module_name, options.log_file, options.verbose)


    #
    #   log programme parameters
    #
    if not options.skip_parameter_logging:
        logger.info(" ".join(sys.argv))




#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Main logic


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
if __name__ == '__main__':

    #
    #   Multiple input files with glob specifications
    #
    import glob
    output_r_file = Popen(["R", "--vanilla", "--slave"], stdin=PIPE)
    #output_r_file = open(options.output_r_file,  "w")

    input_files = []
    #print options.input_files
    for file_group in options.input_files:
        #print "file_group = ", file_group
        for input_glob in file_group:
            #print "input_glob = ", input_glob
            #print glob.glob(input_glob)
            input_files.extend(glob.glob(input_glob))

    counts_per_category = defaultdict(lambda: defaultdict(int))
    for input_file_name in input_files:
        #output_r_file.write("#" + input_file_name + "\n")
        logger.log(MESSAGE, "Parsing " + input_file_name)


        input_file = open(input_file_name)

        histogram_section = False
        cnt_lines = 0
        for line in input_file:

            cnt_lines += 1

            #skip past non histogram sections and first line
            if not len(line):
                continue

            if not histogram_section:
                if line[0:12] == "## HISTOGRAM":
                    fields = input_file.next().rstrip().split('\t')
                    if fields.pop(0) != 'insert_size':
                        raise Exception("First fields not insert_size")
                    histogram_section = True

                continue

            data = line.rstrip().split('\t')
            if len(data) <= 1:
                continue
            data = map(int, data)
            insert_size = data.pop(0)
            for ii, d in enumerate(data):
                if d:
                    counts_per_category[fields[ii]][insert_size] += d



    #
    #   for each category
    #
    for field in "fr_count", "rf_count":
        running_cnt = 0
        cnts_x = [0,1,2]
        cnts_y = [0,1,2]
        if field in counts_per_category:
            for insert_size, cnt in sorted(counts_per_category[field].items()):
                cnts_x.append(insert_size)
                cnts_y.append(cnt)

        #print array_in_r_text_format(field + "_x", cnts_x)
        #print array_in_r_text_format(field + "_y", cnts_y)
        output_r_file.stdin.write(array_in_r_text_format(field + "_x", cnts_x) + "\n")
        output_r_file.stdin.write(array_in_r_text_format(field + "_y", cnts_y) + "\n")
    cmd = """
        pdf("%s", paper="a4r")

        #par(omi = c(0,0,0,0), mar=c(0,0,0,0))
        #par(mfrow=c(2,1), mai=c(0,1,0,0), oma = c(0,0,0,0))

        do_plot <- function(rf_count_x_orig, rf_count_y_orig, fr_count_x_orig, fr_count_y_orig, xranges, bin_size)
        {
            par(omi = c(0,0,0,0), mar=c(0,0,0,0), mai=c(0,1,0,0), oma = c(0,0,0,0))
            layout(matrix(c(1,2), 2, 1, byrow = TRUE), heights = c(2,3),widths = c(1,1))
            rf_count_x = round(rf_count_x_orig / bin_size + bin_size / 2)
            rf_count_y = aggregate(rf_count_y_orig, list(rf_count_x), FUN = sum)$x
            rf_count_x = aggregate(rf_count_x_orig, list(rf_count_x), FUN = mean)$x
            fr_count_x = round(fr_count_x_orig / bin_size + bin_size / 2)
            fr_count_y = aggregate(fr_count_y_orig, list(fr_count_x), FUN = sum)$x
            fr_count_x = aggregate(fr_count_x_orig, list(fr_count_x), FUN = mean)$x

            rf_count_y <- subset(rf_count_y, rf_count_x >= xranges[1] & rf_count_x <= xranges[2])
            rf_count_x <- subset(rf_count_x, rf_count_x >= xranges[1] & rf_count_x <= xranges[2])
            fr_count_y <- subset(fr_count_y, fr_count_x >= xranges[1] & fr_count_x <= xranges[2])
            fr_count_x <- subset(fr_count_x, fr_count_x >= xranges[1] & fr_count_x <= xranges[2])

            plot  (fr_count_x, fr_count_y, type='l', xlim = xranges, xlab = "",
            ylab = "Counts",   col = 'black', omi = c(0,0,0,0), xaxt = 'n')
            points(rf_count_x, rf_count_y, type='l', col = 'blue')

            legend(x = "topleft", fill=c("black", "blue"), bty = "n", border = NULL, legend = c('FR', "RF"))


            par(mai=c(1,1,0,0))
            fr_count_y = cumsum(fr_count_y)
            rf_count_y = cumsum(rf_count_y)
            plot  (fr_count_x, fr_count_y, type='l', xlim = xranges,
            xlab = "Insert sizes", ylab = "Cummulative counts", mar =c(0,0,0,0),
            col = 'black')
            points(rf_count_x, rf_count_y, type='l', col = 'blue')
            legend(x = "topleft", fill=c("black", "blue"), bty = "n", border = NULL, legend = c('FR', "RF"))

            median_index = min(which(fr_count_y > max(fr_count_y) / 2))
            points( c(0, fr_count_x[median_index], fr_count_x[median_index]),
                    c(fr_count_y[median_index], fr_count_y[median_index], 0), col = 'red', type='l')

            text(xranges[1], fr_count_y[median_index]*1.5, pos = 4,
                labels = paste("median = ", fr_count_x[median_index], "\nbp count = ", max(fr_count_y), " pairs between [", xranges[1], " & ", xranges[2], "]"))
        }

        do_plot(rf_count_x, rf_count_y, fr_count_x, fr_count_y, c(20000,50000), 200)
        do_plot(rf_count_x, rf_count_y, fr_count_x, fr_count_y, c(0,1000)     , 1)
        do_plot(rf_count_x, rf_count_y, fr_count_x, fr_count_y, c(0,70000)    , 200)

                              dev.off()
        """ % options.output_pdf_file
    #print cmd
    output_r_file.stdin.write(cmd)
